package com.alpha.gateway.filter;

import com.alpha.core.http.HttpConstants;
import com.alpha.core.http.Response;
import com.alpha.core.utils.JsonUtil;
import com.alpha.gateway.schemas.logs.Log;
import com.alpha.gateway.schemas.logs.LogService;
import com.alpha.gateway.utils.RequestUtil;
import com.alpha.gateway.utils.ResponseUtil;
import lombok.RequiredArgsConstructor;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.*;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * 日志过滤器，拦截请求并记录日志
 *
 * @author chenruwan
 * @since 2022-01-01
 */
@Component
@RequiredArgsConstructor
public class AlphaLoggingFilter implements GlobalFilter, Ordered {
    private final LogService logService;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        if (HttpMethod.GET.name().equalsIgnoreCase(request.getMethodValue())) {
            return chain.filter(exchange.mutate()
                    .response(getResponseDecorator(exchange, request.getURI().getQuery()))
                    .build());
        }

        String contentType = request.getHeaders().getFirst(HttpHeaders.CONTENT_TYPE);
        if (contentType == null || !contentType.contains(MediaType.APPLICATION_JSON_VALUE)) {
            // 当前先严格限制所有请求使用json，遇到图片等特殊场景再特殊处理
            return ResponseUtil.response(exchange, HttpStatus.BAD_REQUEST, "不支持的 Content-Type: " + contentType);
        }

        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        headers.remove(HttpHeaders.CONTENT_LENGTH);

        // 读取请求体
        StringBuilder bodyBuilder = new StringBuilder();
        Mono<String> modifiedBody = ServerRequest.create(exchange, ServerCodecConfigurer.create().getReaders())
                .bodyToMono(String.class)
                .flatMap(body -> {
                    bodyBuilder.append(body);
                    return Mono.just(body);
                });

        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
        return BodyInserters.fromPublisher(modifiedBody, String.class)
                .insert(outputMessage, new BodyInserterContext())
                .then(Mono.defer(() -> chain.filter(exchange.mutate()
                        .request(getRequestDecorator(exchange, headers, outputMessage))
                        .response(getResponseDecorator(exchange, bodyBuilder.toString()))
                        .build())));
    }

    private ServerHttpRequestDecorator getRequestDecorator(ServerWebExchange exchange, HttpHeaders headers,
                                                           CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            @Nonnull
            public HttpHeaders getHeaders() {
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                long contentLength = headers.getContentLength();
                if (contentLength > 0) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, HttpConstants.CHUNKED);
                }
                return httpHeaders;
            }

            @Override
            @Nonnull
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }

    private ServerHttpResponseDecorator getResponseDecorator(ServerWebExchange exchange, String request) {
        ServerHttpResponse response = exchange.getResponse();
        return new ServerHttpResponseDecorator(response) {
            @Override
            @Nonnull
            public Mono<Void> writeWith(@Nonnull Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
                        DataBuffer joinedBuffer = bufferFactory.join(dataBuffers);
                        byte[] context = new byte[joinedBuffer.readableByteCount()]; // probably should reuse buffers
                        joinedBuffer.read(context);
                        joinedBuffer.readPosition(0); // 避免使用DataBufferUtils.release 和 response.bufferFactory().wrap
                        saveLogs(exchange, request, response.getStatusCode(), context);
                        return joinedBuffer;
                    }));
                }
                // if body is not a flux. never got there.
                return super.writeWith(body);
            }
        };
    }

    private void saveLogs(ServerWebExchange exchange, String requestParam, HttpStatus status, byte[] response) {
        ServerHttpRequest request = exchange.getRequest();

        Map<String, Object> claims = JsonUtil.toMap(request.getHeaders().getFirst(HttpConstants.CLAIMS));
        String username = (String) claims.get(HttpConstants.CLAIM_USER_NAME);
        Log log = new Log()
                .setUser((String) claims.get(HttpConstants.CLAIM_USER_ID))
                .setClient(request.getHeaders().getFirst(HttpConstants.HEADER_CLIENT_NAME))
                .setUrl(request.getPath().value())
                .setMethod(request.getMethodValue())
                .setStatus(status == null ? null : status.value())
                .setMessage(status == null ? null : status.getReasonPhrase())
                .setRemoteAddress(RequestUtil.getRemoteIp(request))
                .setRequest(requestParam)
                .setCreatedBy(username)
                .setUpdatedBy(username)
                .build();

        if (HttpMethod.GET.matches(request.getMethodValue())) {
            Response<?> resp = JsonUtil.fromString(new String(response, StandardCharsets.UTF_8), Response.class);
            if (resp != null) {
                resp.setData(null);
                log.setResponse(JsonUtil.toString(resp));
            } else {
                log.setResponse(null);
            }
        } else {
            log.setResponse(new String(response, StandardCharsets.UTF_8));
        }
        logService.save(log);
    }

    @Override
    public int getOrder() {
        return -4;
    }
}
